# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: inkyu

import logging
import math
from collections import namedtuple

from coin.strategy.hk_hard_arb.util import PassAgg, PassAggExchangeMap

OrderSizes = namedtuple('OrderSizes', ['timestamp', 'buy_qty', 'sell_qty'])


class BalanceManager(object):
  def __init__(self, *, order_delay, stable_period, exchange_side=PassAgg.PASSIVE, logger=None):
    self._base_balance = None
    self._last_stable_balance = None

    self._offset_qty = 0.
    self._timestamp_last_change = 0
    self._timestamp_untrustable = None

    self._order_delay = order_delay
    self._stable_period = stable_period
    self._exchange_side = exchange_side
    self._logger = logger or logging.getLogger(__name__)

    self._stable = True

  @property
  def stable(self):
    return self._stable

  def add_offset_qty(self, *, offset_qty, timestamp):
    self._offset_qty += offset_qty
    self._timestamp_last_change = timestamp
    self.untrust(timestamp=timestamp)
    self._stable = False

  def reset_offset_qty(self, *, timestamp):
    self._offset_qty = 0.
    self._timestamp_last_change = timestamp

  def get_balance(self, *, use_snapshot=False):
    if use_snapshot:
      return self._base_balance

    if self._last_stable_balance is not None:
      return self._offset_qty + self._last_stable_balance
    else:
      return None

  def update(self, base_balance, timestamp, *, on_stabilize=None):
    if (self._base_balance is None or not math.isclose(self._base_balance, base_balance)):
      if self._timestamp_last_change and self._offset_qty == 0.:
        self._logger.warn('Unexpected balance update: %s => %s', self._base_balance, base_balance)
      if self._timestamp_last_change:
        self._stable = False
      self._timestamp_last_change = timestamp

    if self.is_stable(timestamp):
      if self._offset_qty != 0:
        self._logger.info(
            'Stabilized %s: (%f, %f) => %f' %
            (self._exchange_side, self._last_stable_balance, self._offset_qty, base_balance))
        if on_stabilize:
          on_stabilize(base_balance)
      self._stable = True
      self._last_stable_balance = base_balance
      self.reset_offset_qty(timestamp=timestamp)
    if self._last_stable_balance is None:
      self._last_stable_balance = base_balance

    self._base_balance = base_balance

  def untrust(self, timestamp):
    self._timestamp_untrustable = timestamp

  def is_trustable(self, timestamp):
    if self._timestamp_untrustable is not None:
      if timestamp > self._timestamp_untrustable + self._order_delay:
        self._timestamp_untrustable = timestamp
        return True
      else:
        return False
    else:
      return True

  def is_stable(self, timestamp, verbose=False):
    if timestamp > self._timestamp_last_change + self._stable_period:
      return True
    else:
      return False


class SingleProductSizeCalculator(object):
  def __init__(self,
               *,
               base_reserve,
               lot_size,
               delay_aggressive_order=5e9,
               stable_period=3e9,
               logger=None):
    self._balance_aggressive = BalanceManager(exchange_side=PassAgg.AGGRESSIVE,
                                              order_delay=delay_aggressive_order,
                                              stable_period=delay_aggressive_order,
                                              logger=logger)
    self._balance_passive = BalanceManager(exchange_side=PassAgg.PASSIVE,
                                           order_delay=0e9,
                                           stable_period=stable_period,
                                           logger=logger)

    self._base_reserve = base_reserve  # TODO(inkyu): Merge with FungibleArbBalance
    self._lot_size = lot_size
    self._timestamp = None
    self._balance = None
    self._ready = False
    self._logger = logger or logging.getLogger(__name__)

    self.on_stabilize_aggressive = None

  def is_stable(self, timestamp, verbose=False):
    if verbose:
      self._logger.info('%s, %s', self._balance_aggressive.stable, self._balance_passive.stable)

    return (self._balance_aggressive.stable and self._balance_passive.stable)

  @property
  def timestamp(self):
    return self._timestamp

  @property
  def ready(self):
    return self._ready

  def get_imbalance(self, *, use_snapshot=False):
    base_total = 0.
    balance_aggressive = self._balance_aggressive.get_balance(use_snapshot=use_snapshot)
    if balance_aggressive is None:
      return None
    base_total += balance_aggressive
    balance_passive = self._balance_passive.get_balance(use_snapshot=use_snapshot)
    if balance_passive is None:
      return None
    base_total += balance_passive
    return base_total - self._base_reserve

  def update(self, timestamp, balance):
    self._timestamp = timestamp
    self._balance = balance
    if len(balance.base_balances) > 1:
      base_balances_passagg = PassAggExchangeMap._make(balance.base_balances)
      self._balance_aggressive.update(base_balances_passagg.aggressive,
                                      timestamp,
                                      on_stabilize=self.on_stabilize_aggressive)
      self._balance_passive.update(base_balances_passagg.passive, timestamp)
    else:
      self._balance_passive.update(balance.base_balances[0], timestamp)
    self._ready = True

  def on_passive_fill_buy(self, *, qty_buy, timestamp):
    assert qty_buy > 0, qty_buy
    self._balance_passive.add_offset_qty(offset_qty=qty_buy, timestamp=timestamp)

  def on_passive_fill_sell(self, *, qty_sell, timestamp):
    assert qty_sell > 0, qty_sell
    self._balance_passive.add_offset_qty(offset_qty=-qty_sell, timestamp=timestamp)

  def on_aggressive_fire_buy(self, *, qty_buy, timestamp):
    assert qty_buy > 0, qty_buy
    self._balance_aggressive.untrust(timestamp)
    self._balance_aggressive.add_offset_qty(offset_qty=qty_buy, timestamp=timestamp)

  def on_aggressive_fire_sell(self, *, qty_sell, timestamp):
    assert qty_sell > 0, qty_sell
    self._balance_aggressive.untrust(timestamp)
    self._balance_aggressive.add_offset_qty(offset_qty=-qty_sell, timestamp=timestamp)

  def calculate_size_passive(self):
    imbalance = self.get_imbalance()
    if imbalance is None:  # Not ready
      return None
    buy_size = self._lot_size
    sell_size = self._lot_size

    if self._balance_aggressive.is_trustable(self._timestamp):
      if imbalance < 0:
        buy_size += min(self._lot_size, -imbalance)
      elif imbalance > 0:
        sell_size += min(self._lot_size, imbalance)

    return OrderSizes(timestamp=self.timestamp, buy_qty=buy_size, sell_qty=sell_size)

  def calculate_size_aggressive(self):
    imbalance = self.get_imbalance()
    imbalance_snapshot = self.get_imbalance(use_snapshot=True)
    if imbalance is None or imbalance_snapshot is None:  # Not ready
      return None
    buy_size = max(-imbalance, 0)
    sell_size = max(imbalance, 0)

    balance_aggressive = self._balance_aggressive.get_balance()
    sell_size = min(sell_size, balance_aggressive)

    # For the safety
    buy_size_max = max(self._lot_size * 2.0 - imbalance_snapshot, 0)
    sell_size_max = max(self._lot_size * 2.0 + imbalance_snapshot, 0)
    return OrderSizes(timestamp=self.timestamp,
                      buy_qty=min(buy_size, buy_size_max),
                      sell_qty=min(sell_size, sell_size_max))
